# 请求结果对比
import os
import json
import argparse

from numpy import newaxis

try:
    import numpy as np
except:
    os.system("pip install --upgrade " + "numpy")

try:
    import pandas as pd
except:
    os.system("pip install --upgrade " + "pandas")

try:
    import yaml
except:
    os.system("pip install --upgrade " + "pyyaml")

try:
    from sklearn.metrics import mean_squared_error, r2_score, mean_absolute_error
except:
    os.system("pip install --upgrade " + "scikit-learn")

def parse_args():
    parser = argparse.ArgumentParser("集中化IT智能决策平台返回结果对比脚本")

    # 配置文件
    parser.add_argument('-c',
                        '--config',
                        type=str,
                        default='./configs/compare_config.yaml',
                        help='配置文件的路径')
    
    args = parser.parse_args()
    return args

def process(response_file:str, compare_file:str, out_path:str):

    # 首先创建文件保存路径
    if not os.path.exists(out_path):
        os.makedirs(out_path)

    with open(response_file, 'r') as f:
        response_data = json.load(f)
    f.close()

    with open(compare_file, 'r') as f:
        compare_data = json.load(f)
    f.close()

    result_array =[]
    result_array.append(['system', None, response_data['data']['system'], compare_data['data']['system'], compare_data['data']['system']- response_data['data']['system']])

    response_array=[]
    response_array.append(['system', 'system', response_data['data']['system']])
    for dim in response_data['data'].keys():
        if dim == 'system':
            continue
        else:
            dim_scores = response_data['data'][dim]
            for dim_res in dim_scores:
                dim_res_id = dim_res['id']
                dim_res_score = dim_res['score']
                response_array.append([dim, dim_res_id, dim_res_score])
    response_df = pd.DataFrame(response_array, columns=['维度', '资源ID', '测试得分'])
    # print(response_df)

    compare_array=[]
    compare_array.append(['system', 'system', compare_data['data']['system']])
    for dim in compare_data['data'].keys():
        if dim == 'system':
            continue
        else:
            dim_scores = compare_data['data'][dim]
            for dim_res in dim_scores:
                dim_res_id = dim_res['id']
                dim_res_score = dim_res['score']
                compare_array.append([dim, dim_res_id, dim_res_score])
    compare_df = pd.DataFrame(compare_array, columns=['维度', '资源ID', 'GT得分'])
    # print(compare_df)

    result_df = pd.merge(response_df, compare_df, how='inner', on=['维度', '资源ID'])
    MSE = mean_squared_error(result_df['测试得分'], result_df['GT得分'])
    RMSE = np.sqrt(MSE)
    MAE = mean_absolute_error(result_df['测试得分'], result_df['GT得分'])
    R2_SCORE = r2_score(result_df['测试得分'], result_df['GT得分'])

    # 保存对比结果
    result_df.to_csv(os.path.join(out_path, 'compare_result.csv'), index=False)
    with open(os.path.join(out_path, 'compare_result.txt'), 'w') as f:
        for key, value in zip(['RMSE', 'MAE', 'R2'], [RMSE, MAE, R2_SCORE]):
            f.write(key+": " + str(value) + "\n")
    f.close()

    return MSE, RMSE, R2_SCORE
    
    
def main():
    args = parse_args()
    with open(args.config, 'r') as fc:
        config = yaml.load(fc, Loader=yaml.FullLoader)
    fc.close()
    
    # 创建存放对比结果的文件夹路径
    out_dir = config['output_dir']
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)
    
    mae_list = []
    rmse_list = []
    r2_list = []
    for item in config['datas']:
        index = item['index']
        response_file = item['response_file']
        compare_file = item['gt_file']
        file_folder = os.path.join(out_dir, str(index))
        rmse, mae, r2 = process(response_file, compare_file, file_folder)
        mae_list.append(mae)
        rmse_list.append(rmse)
        r2_list.append(r2)
    
    with open(os.path.join(out_dir, 'compare_result.txt'), 'w') as f:
        for key, value in zip(['RMSE', 'MAE' ,'R2'], [np.mean(rmse_list), np.mean(mae_list), np.mean(r2_list)]):
            f.write(key+": " + str(value) + "\n")
    f.close()

    # for item in config:
    #     process(item['response_file'], item['compare_file'])

if __name__ == "__main__":
    main()

